package top.minuy.structure.map;

import top.minuy.structure.linear.array.Array;

/**
 * 基于红黑树实现映射，具体注释看红黑树
 *
 * @param <K> 键，需要可比较
 * @param <V> 值
 * @author Minuy
 * @time 23:04
 * @date 2021/11/6
 */
public class RBTreeMap<K extends Comparable<K>, V> implements Map<K, V> {

    private final static boolean RED = true;
    private final static boolean BLACK = false;

    class Node {
        // 键
        public K key;
        // 值
        public V val;
        // 右子树
        public Node right;
        // 左子树
        public Node left;

        public boolean color;

        /**
         * 初始化一个节点
         *
         * @param key 节点的键
         * @param val 节点的值
         */
        public Node(K key, V val) {
            this.key = key;
            this.val = val;
            this.right = null;
            this.left = null;
            this.color = RED;
        }
    }

    /**
     * 根节点
     */
    private Node root;
    /**
     * 大小
     */
    private int size;

    public RBTreeMap() {
        root = null;
        size = 0;
    }

    private Node leftRotate(Node node) {
        Node x = node.right;

        node.right = x.left;
        x.left = node;

        x.color = node.color;
        node.color = RED;

        return x;
    }

    private Node rightRotate(Node node) {
        Node x = node.left;

        node.left = x.right;
        x.right = node;

        x.color = node.color;
        node.color = RED;

        return x;
    }

    private void flipColors(Node node) {
        node.color = RED;
        node.right.color = BLACK;
        node.left.color = BLACK;
    }

    private boolean isRed(Node node) {
        if (node == null) {
            return false;
        }
        return node.color == RED;
    }

    /**
     * 向字典中添加键值对
     *
     * @param key 键
     * @param val 值
     */
    @Override
    public void add(K key, V val) {
        root = add(root, key, val);
    }

    public Node add(Node node, K key, V val) {
        if (node == null) {
            size++;
            return new Node(key, val);
        }

        if (key.compareTo(node.key) < 0) {
            node.left = add(node.left, key, val);
        } else if (key.compareTo(node.key) > 0) {
            node.right = add(node.right, key, val);
        } else {  // key.compareTo(node.key) == 0
            node.val = val;
        }

        if (!isRed(node.left) && isRed(node.right)) {
            node = leftRotate(node);
        }

        if (isRed(node.left) && isRed(node.left.left)) {
            node = rightRotate(node);
        }

        if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }

        return node;
    }

    /**
     * 获取这棵树中的最小值
     *
     * @param node 根节点
     * @return 最小值所在的节点
     */
    private Node getMinimum(Node node) {
        if (node == null) {
            return null;
        }
        if (node.left == null) {
            return node;
        } else {
            return getMinimum(node.left);
        }
    }

    /**
     * 删除一个键值对
     *
     * @param key 要被删除的键
     * @return 返回被删除键的值，如若键不存在，返回空
     */
    @Override
    public V remove(K key) {
        Node node = getNode(root, key);
        if (node == null) {
            return null;
        } else {
            root = remove(root, key);
            return node.val;
        }
    }

    private Node remove(Node node, K key) {
        if (node == null) {
            return null;
        }
        if (node.key.compareTo(key) > 0) {
            node.left = remove(node.left, key);
            return node;
        } else if (node.key.compareTo(key) < 0) {
            node.right = remove(node.right, key);
            return node;
        } else { // node.key.compareTo(key) == 0
            if (node.left == null) {
                Node t = node.right;
                node.right = null;
                size--;
                return t;
            }

            if (node.right == null) {
                Node t = node.left;
                node.left = null;
                size--;
                return t;
            }

            Node minimumNode = getMinimum(node.right);

            node.right = remove(node.right, minimumNode.key);

            minimumNode.left = node.left;
            minimumNode.right = node.right;

            node.right = node.left = null;

            return minimumNode;
        }
    }

    /**
     * 根据key获取节点
     *
     * @param root 根节点
     * @param key  要查找的key
     * @return 对应的节点或者null
     */
    private Node getNode(Node root, K key) {
        if (root == null) {
            return null;
        }

        if (root.key.compareTo(key) > 0) {
            return getNode(root.left, key);
        } else if (root.key.compareTo(key) < 0) {
            return getNode(root.right, key);
        } else { // root.key.compareTo(key) == 0
            return root;
        }
    }

    /**
     * 查询键是否存在
     *
     * @param key 要查询的键
     * @return 此键是否存在
     */
    @Override
    public boolean contains(K key) {
        return getNode(root, key) != null;
    }

    /**
     * 获取一个键的值
     *
     * @param key 键
     * @return 值
     */
    @Override
    public V get(K key) {
        // 当根节点为空的时候返回空，而不是空指针异常
        Node t = getNode(root, key);
        return t == null ? null : t.val;
    }

    /**
     * 修改一个键值对
     *
     * @param key    键
     * @param newVal 值
     */
    @Override
    public void set(K key, V newVal) {
        Node node = getNode(root, key);
        if (node == null) {
            throw new IllegalArgumentException("Set failed. The \"" + key.toString() + "\" is not exists.");
        } else {
            node.val = newVal;
        }
    }

    /**
     * 获取字典的大小
     *
     * @return 字典的大小
     */
    @Override
    public int getSize() {
        return size;
    }

    /**
     * 判断字典是否为空
     *
     * @return 字典是否为空
     */
    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * 判断这棵树是不是二分搜索树
     *
     * @return 是/不是
     */
    public boolean isBST() {
        Array<K> keys = new Array<>(size);
        inOrder(root, keys);
        for (int i = 1; i < keys.getSize(); i++) {
            if (keys.get(i - 1).compareTo(keys.get(i)) > 0) {
                return false;
            }
        }
        return true;
    }

    /**
     * 为isBST()服务，中序遍历，放在Array中
     *
     * @param node  遍历开始的结点
     * @param array 数组
     */
    private void inOrder(Node node, Array<K> array) {
        if (node == null) {
            return;
        }
        try {
            inOrder(node.left, array);
            array.addLast(node.key);
            inOrder(node.right, array);
        } catch (StackOverflowError e) {
            System.out.println("Index : " + array.getSize() + " Error : " + e.toString());
        }
    }
}
